# -*- coding: utf-8 -*-
import os
import pandas as pd
import time
from datetime import datetime
import matplotlib.pyplot as plt
import matplotlib.ticker as mick
import joblib as job
from sklearn.metrics import mean_absolute_error
from rich.traceback import install
from warnings import filterwarnings
from utils.log import Logger
from utils.common import preprocessing, TIME_FORMAT, LOG_TIME_FORMAT

install()  # 报错警告信息优化
filterwarnings('ignore', module='sklearn')  # 忽略警告信息
# 设置日志文件名、日志级别
logfile_name = 'predict-' + datetime.now().strftime(LOG_TIME_FORMAT)
logfile = Logger(root_path='../', log_name=logfile_name).get_logger(log_dir='log/predict')
# 读取数据并预处理
power_load_data = preprocessing('../data/test.csv')
time_load_dict = power_load_data.set_index('time')['power_load'].to_dict()


def feature_extract(pred_time: list, time_load: dict) -> pd.DataFrame:
    """
    ## feature_extract\n
    Based on the predicted time and historical load, process feature data
    :param pred_time: Predicted time
    :param time_load: Historical load before the predicted time
    :return: Feature engineering processed data
    """
    hour_cols = [f'hour_{i:02d}' for i in range(24)]
    month_cols = [f'month_{i:02d}' for i in range(1, 13)]
    feature_cols = hour_cols + month_cols + ['前1小时负荷', '前2小时负荷', '前3小时负荷', 'yesterday_load']
    # 小时特征
    pred_hour = pred_time[11:13]
    hour_list = []
    # 小时转为编码
    for i in range(24):
        if pred_hour == feature_cols[i][5:7]:
            hour_list.append(1)
        else:
            hour_list.append(0)
    # 月份特征
    pred_month = pred_time[5:7]
    month_list = []
    # 月份转为编码
    for i in range(24, 36):
        if pred_month == feature_cols[i][6:8]:
            month_list.append(1)
        else:
            month_list.append(0)
    # 前1小时、2小时、3小时、昨日的负荷
    last_hour_day_list = []
    for i in ['1h', '2h', '3h', '1d']:
        last_hour_day_time = (pd.to_datetime(pred_time) - pd.to_timedelta(i)).strftime(TIME_FORMAT)
        last_hour_day_load = time_load.get(last_hour_day_time, 500)
        last_hour_day_list.append(last_hour_day_load)
    feature_list = [hour_list + month_list + last_hour_day_list]
    feature_data = pd.DataFrame(feature_list, columns=feature_cols)
    return feature_data


def predict(data: pd.DataFrame, time_load: dict) -> pd.DataFrame:
    """
    ## predict\n
    Predict the load at the specified time
    :param data: Power load data
    :param time_load: Load before the predicted time
    :return: Evaluated feature engineering precessing data
    """
    # 模型预测, 确定要预测的时间段 (2015-08-01 00:00:00 及以后的时间)
    predict_times = [t for t in data['time'] if t >= '2015-08-01 00:00:00']
    evaluate_list = []
    start_time = time.time()
    logfile.info('Start predict time')
    for predict_time in predict_times:
        # 为了模拟真实预测, 把要预测的时间以及以后的负荷都掩盖掉
        # 新建一个数据字典, 只保存预测时间以前的数据字典
        time_load_dict_masked = {k: v for k, v in time_load.items() if k < predict_time}
        # 预测负荷, 解析特征
        processed_data = feature_extract(predict_time, time_load_dict_masked)
        # 利用加载的模型预测
        y_predict = xgb_model.predict(processed_data)
        # 保存预测时间对应的真实负荷
        true_value = time_load.get(predict_time)
        # 结果保存到 evaluate_list, 三个元素分别是预测时间、真实负荷、预测负荷
        evaluate_list.append([predict_time, true_value, y_predict[0]])
    # evaluate_list 转为 DataFrame
    evaluate_data = pd.DataFrame(evaluate_list, columns=['预测时间', '真实负荷', '预测负荷'])
    # 预测结果评价, 计算预测结果与真实结果的 MAE
    mae = mean_absolute_error(evaluate_data['真实负荷'], evaluate_data['预测负荷'])
    logfile.info(f'MAE: {mae}')
    end_time = time.time()
    logfile.info('End predict time')
    logfile.info(f'Processing time: {end_time - start_time}s')
    return evaluate_data


def data_analyze(data: pd.DataFrame) -> None:
    """
    ## data_analyze\n
    Result analysis, draw a predicted time actual load line chart, predicted time predicted load line chart
    :param data: Predicted results `['预测时间', '真实负荷', '预测负荷']`
    :return: None
    """
    fig = plt.figure(figsize=(40, 20))
    ax = fig.add_subplot()
    ax.plot(data['预测时间'], data['真实负荷'], label='actual load')
    ax.plot(data['预测时间'], data['预测负荷'], label='predicted load')
    ax.set_title('Comparison of actual load and predicted load')
    ax.set_xlabel('predicted time')
    ax.set_ylabel('load')
    # 横坐标时间若不处理太过密集, 这里调大时间展示的间隔
    ax.xaxis.set_major_locator(mick.MultipleLocator(50))
    # 时间展示时旋转45度
    plt.xticks(rotation=45)
    plt.legend()
    plt.savefig('../images/comparison_actual_predict_load.png')
    logfile.info('`comparison_actual_predict_load.png` saved!')


if __name__ == '__main__':
    # 如果 `xgb_power_load.pth` 模型文件存在则加载模型
    try:
        if os.path.exists('../model/xgb_power_load.pth'):
            xgb_model = job.load('../model/xgb_power_load.pth')
        else:
            logfile.error('`xgb_power_load.pth` not exists!')
            exit()
    except Exception as e:
        logfile.error(e)
    # 如果 `comparison_actual_predict_load.png` 图片存在则不运行 `time_predict` 函数
    try:
        if not os.path.exists('../images/comparison_actual_predict_load.png'):
            predict_data = predict(power_load_data, time_load_dict)
            data_analyze(predict_data)
        else:
            logfile.warning('`comparison_actual_predict_load.png` already exists!')
    except Exception as e:
        logfile.error(e)
